# -*- coding: utf-8 -*-
"""
# @FileName:     sample_ddim.py
# @AuthorName:   Sanqi Lu (Lingwei Dang)
# @Institution:  SCUT, Guangzhou, China
# @EmailAddress: lenvondang@163.com
# @CreateTime:   2024/12/22 17:11
"""
import sys
sys.path.append("")
import torch
import torch.nn as nn
from ddim.models.ddim import DDIM
import cv2
import numpy as np
import einops
from tqdm import tqdm

from ddpm.cfgs.ddpm_cfgs import unet_res_cfg
from ddpm.datas.get_dataloader import get_dataloader
from ddpm.models.net import build_network

batch_size = 512
n_epochs = 100

def sample_imgs(ddim,
                net,
                n_sample=81,
                device='cuda',
                simple_var=True):
    model_path = f"outs/last_099.pth"
    net.load_state_dict(torch.load(model_path))
    net = net.to(device)
    net = net.eval()
    # for ddim_step in [1000, 500, 300, 200, 100, 50, 20, 10, 5]:
    for ddim_step in [100, 50, 20, 10, 5]:
        img_save_path = f"outs/ddim_step_{ddim_step:03d}.png"
        with torch.no_grad():
            shape = (n_sample, 1, 28, 28)  # 1, 3, 28, 28
            imgs = ddim.sample_backward(shape,
                                        net,
                                        device=device,
                                        simple_var=simple_var,
                                        ddim_step=ddim_step, eta=1
                                        ).detach().cpu()
            imgs = (imgs + 1) / 2 * 255
            imgs = imgs.clamp(0, 255)
            imgs = einops.rearrange(imgs,
                                    '(b1 b2) c h w -> (b1 h) (b2 w) c',
                                    b1=int(n_sample**0.5))

            imgs = imgs.numpy().astype(np.uint8)

            cv2.imwrite(img_save_path, imgs)

if __name__ == '__main__':
    n_steps = 1000
    config_id = 4
    device = 'mps:0'
    config = unet_res_cfg
    net = build_network(config, n_steps)
    ddpm = DDIM(device, n_steps)

    sample_imgs(ddpm, net, device=device)
